import sys
sys.path.append("/Users/yishanli/python/myProject/python-demo/test/")

import backtrader as bt
import pandas as pd
import datetime
from tushares.TuShareTool import TuShareTool  # 导入用于获取股票数据的TuShareTool
# 初始化 TuShareTool
tushare = TuShareTool()

# 获取股票数据
def get_data(code='600483.SH', start_date='2022-10-01', end_date='2023-10-01'):
    # 使用TuShareTool获取指定股票代码和日期范围内的股票数据
    df = tushare.get_stock_data(code, start_date, end_date)

    # 调整数据格式，将日期列转换为Datetime类型并设置为索引
    df['Datetime'] = pd.to_datetime(df['trade_date'])
    df.set_index("Datetime", inplace=True)

    # 添加一个名为'openinterest'的列，值都设置为0
    df['openinterest'] = 0

    # 重命名'vol'列为'volume'
    df.rename(columns={'vol': 'volume'}, inplace=True)

    # 选择并返回特定的列，同时反转数据，以便按照时间升序排列
    df = df[['open', 'high', 'low', 'close', 'volume', 'openinterest']]
    return df.iloc[::-1]  # 反转数据，按时间升序排列


# 恒瑞医药
data1 = get_data('600276.SH','20200101','20211015')
# 贵州茅台
data2 = get_data('600519.SH','20200101','20211015')
# 海天味业
data3 = get_data('603288.SH','20200101','20211015')

class MyStrategy(bt.Strategy):

    def __init__(self):
        # 打印数据集和数据集对应的名称
        print("-------------self.datas-------------")
        print(self.datas)
        print("-------------self.data-------------")
        print(self.data._name, self.data) # 返回第一个导入的数据表格，缩写形式
        print("-------------self.data0-------------")
        print(self.data0._name, self.data0) # 返回第一个导入的数据表格，缩写形式
        print("-------------self.datas[0]-------------")
        print(self.datas[0]._name, self.datas[0]) # 返回第一个导入的数据表格，常规形式
        print("-------------self.datas[1]-------------")
        print(self.datas[1]._name, self.datas[1]) # 返回第二个导入的数据表格，常规形式
        print("-------------self.datas[-1]-------------")
        print(self.datas[-1]._name, self.datas[-1]) # 返回最后一个导入的数据表格
        print("-------------self.datas[-2]-------------")
        print(self.datas[-2]._name, self.datas[-2]) # 返回倒数第二个导入的数据表格

class MyStrategy01(bt.Strategy):
    def __init__(self):
        print("--------- 打印 self 策略本身的 lines ----------")
        print(self.lines.getlinealiases())
        print("--------- 打印 self.datas 第一个数据表格的 lines ----------")
        print(self.datas[0].lines.getlinealiases())
        # 计算第一个数据集的s收盘价的20日均线，返回一个 Data feed
        self.sma = bt.indicators.SimpleMovingAverage(self.datas[0].close, period=20)
        print("--------- 打印 indicators 对象的 lines ----------")
        print(self.sma.lines.getlinealiases())
        print("---------- 直接打印 indicators 对象的所有 lines -------------")
        print(self.sma.lines)
        print("---------- 直接打印 indicators 对象的第一条 lines -------------")
        print(self.sma.lines[0])
        
    def next(self):
        print('验证索引位置为 6 的线是不是 datetime')
        # datetime 线中的时间点存的是数字形式的时间，可以通过 bt.num2date() 方法将其转为“xxxx-xx-xx xx:xx:xx”这种形式
        print(bt.num2date(self.datas[0].lines[6][0]))

class MyStrategy02(bt.Strategy):
    def __init__(self):
        self.count = 0
        print("-------------init 中的索引位置-------------------")
        print("0 索引:",'datetime',self.data1.lines.datetime.date(0),'close',self.data1.lines.close[0])
        print("-1 索引：",'datetime',self.data1.lines.datetime.date(-1),'close', self.data1.lines.close[-1])
        print("-2 索引",'datetime', self.data1.lines.datetime.date(-2),'close', self.data1.lines.close[-2])
        print("1 索引：",'datetime',self.data1.lines.datetime.date(1),'close', self.data1.lines.close[1])
        print("2 索引",'datetime', self.data1.lines.datetime.date(2),'close', self.data1.lines.close[2])
        # 通过 get() 切片时，如果是从 ago=0 开始取，不会返回数据，从其他索引位置开始取，能返回数据
        print("从 0 开始往前取3天的收盘价：", self.data1.lines.close.get(ago=0, size=3))
        print("从-1开始往前取3天的收盘价：", self.data1.lines.close.get(ago=-1, size=3))
        print("从-2开始往前取3天的收盘价：", self.data1.lines.close.get(ago=-2, size=3))
        print("line的总长度：", self.data1.buflen())
    def next(self):
        print(f"------------- next 的第{self.count+1}次循环 --------------")

        print("已处理的数据点：", len(self.data1))
        print("line的总长度：", self.data0.buflen())
        self.count += 1




if __name__ == '__main__':
    fromdate = datetime.datetime(2020, 1, 1)  # 设置回测开始日期
    todate = datetime.datetime(2021, 10, 15)  # 设置回测结束日期
    cerebro = bt.Cerebro()  # 创建Backtrader大脑

    # 添加 600276.SH 的行情数据
    data1 = bt.feeds.PandasData(dataname=data1,fromdate=fromdate,todate=todate)  # 创建数据源，使用Pandas数据
    cerebro.adddata(data1, "600276.SH")  # 添加数据到大脑，并命名为"pingan"

    # 添加 600519.SH 的行情数据
    data2 = bt.feeds.PandasData(dataname=data2, fromdate=fromdate, todate=todate)
    cerebro.adddata(data2, name='600519.SH')

    # 添加 603288.SH 的行情数据
    data3 = bt.feeds.PandasData(dataname=data3, fromdate=fromdate, todate=todate)
    cerebro.adddata(data3, name='603288.SH')

    cerebro.addstrategy(MyStrategy02)  # 向大脑添加自定义策略

    # 执行回测
    cerebro.run()



